#include <iostream>
#include "util/distance.hpp"
#include "bruteforce.hpp"
#include "util/time_memory.hpp"
#include <vector>
#include <cstring>
#include <unordered_set>
#include <queue>

#undef NDEBUG

#include <cassert>

namespace GraphGradient {

    class KNNGraph {
    public:
        KNNGraph(int vec_dim, int max_element, int k_graph) : vec_dim_(vec_dim), max_element_(max_element),
                                                              k_graph_(k_graph) {
            bfs_ = BruteForceSearch(vec_dim, max_element);
            graph_ = std::vector<std::unordered_set<int>>(max_element);
        }

        KNNGraph() {}

        void build(float *item);

        void test();

        std::vector<std::unordered_set<int>> graph_;

        ~KNNGraph() {
            graph_.clear();
        }

    private:
        BruteForceSearch bfs_;

        int max_element_;
        float *data_;
        int vec_dim_;
        int k_graph_;
        int report_every_ = 5000;
    };

    void KNNGraph::build(float *item) {
        TimeRecord single_insert_record;
        bfs_.insert(item);
        data_ = item;
        for (int i = 0; i < max_element_; i++) {
            graph_[i] = bfs_.search_knn(item + i * vec_dim_, k_graph_);
            if (i % report_every_ == 0) {
                std::cout << "KNNGraph base-base-gnd: " << i / (0.01 * max_element_) << " %, "
                          << 1e-6 * single_insert_record.get_elapsed_time_micro() << " s/iter" << " Mem: "
                          << get_current_RSS() / 1000000 << " Mb \n";
                single_insert_record.reset();
            }
        }
        for (int i = 0; i < max_element_; i++) {
            std::unordered_set<int> neighbors = graph_[i];
            for (std::unordered_set<int>::iterator iter = neighbors.begin(); iter != neighbors.end(); ++iter) {
                int nearest_neighbor = *iter;
                if (graph_[nearest_neighbor].find(i) == graph_[nearest_neighbor].end()) {
                    graph_[nearest_neighbor].insert(i);
                }
            }
        }
    }

    void KNNGraph::test() {
        long n_edge = 0;
        for (int i = 0; i < max_element_; i++) {
            n_edge += graph_[i].size();
            std::unordered_set<int> neighbors = graph_[i];
            for (std::unordered_set<int>::iterator iter = neighbors.begin(); iter != neighbors.end(); ++iter) {
                int nearest_neighbor = *iter;
                assert(graph_[nearest_neighbor].find(i) != graph_[nearest_neighbor].end());
            }
            assert((int) graph_[i].size() >= k_graph_);
        }
        assert(n_edge % 2 == 0);
        std::printf("test pass!\n");
    }


}